from utils.config_load import load_prompt
from utils.llm import llm_api
from schema import db_connect
from utils import logger


def sql_standard(sql: str) -> str:
    """
    标准化 SQL
    :param sql: LLM 生成的 SQL
    :return: 标准化后的 SQL
    """
    return sql.replace('\n', ' ').replace('`', '').replace('sql', '').replace(';', '').strip()


def generate_sql(question: str, task='text2sql', case_qa=None) -> str:
    """
    通过 API 调用 LLM，生成 SQL
    :param question: 用户输入
    :param task: 任务类型，默认执行 text-to-sql 任务
    :param case_qa: 问答案例
    :return: LLM 回复
    """
    # 调用 LLM 生成 SQL
    sql = llm_api(load_prompt(user_input=question, task=task, case_qa=case_qa))
    if sql == "None":
        logger.info(f"fail 【generate sql】 无法生成 SQL")
        return "None"
    else:
        # 标准化 SQL
        return sql_standard(sql)


def rewrite_sql(question: str, sql: str, error_info: str, case_qa=None, task='rewriteSql') -> str:
    """
    根据用户提问和数据库信息，重新生成 SQL
    :param question: 用户输入
    :param sql: 报错 SQL
    :param task: 任务类型，默认执行 rewrite sql 任务
    :param error_info: 执行 SQL 报错详情
    :param case_qa: 问答案例
    :return: LLM 回复
    """
    # 调用 LLM 生成 SQL
    sql = llm_api(load_prompt(user_input=question, task=task, case_qa=case_qa, sql=sql, error_info=error_info))
    # 标准化 SQL
    return sql_standard(sql)


def fix_sql(sql: str, question: str, case_qa=None, iterate=3) -> tuple:
    """
    检验 SQL 是否正确。若执行时报错，默认至多重新生成3次
    :param sql: 待检验的 SQL
    :param question: 用户提问
    :param case_qa: 问答案例
    :param iterate: 迭代次数，默认为3
    :return: 校正后的 SQL 及数据库查询结果
    """
    # 记录 SQL 执行结果
    error_info = None
    sql_result = None
    # 迭代生成 SQL
    for i in range(0, iterate):
        execute_result = db_connect.execute_sql(sql)
        if list(execute_result.keys())[0] == 'fail':  # 若 SQL 执行报错，则重新生成
            error_info = str(list(execute_result.values())[0])
            logger.info(f"{i+1} time 【rewrite sql】 {error_info} 【sql】 {sql}")
            sql = rewrite_sql(question, sql, error_info, case_qa)
        else:
            error_info = None
            sql_result = str(list(execute_result.values())[0])
            break
    if error_info:
        logger.info(f"fail 【generate sql】 无法生成 SQL")
        return None, "目前暂无相关数据"
    else:
        logger.info(f"success 【generate sql】 {sql}")
        return sql, sql_result


if __name__ == '__main__':
    test_sql = """
    ```sql
    SELECT column FROM table;
    ```
    """
    print(sql_standard(test_sql))
